package co.edu.udistrital.lucene;

import java.io.FileOutputStream;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.Set;
import java.util.TreeMap;
import java.util.TreeSet;

import org.apache.lucene.document.Document;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.index.Term;
import org.apache.lucene.index.TermDocs;
import org.apache.lucene.index.TermEnum;
import org.apache.lucene.index.TermFreqVector;
import org.apache.lucene.store.Directory;

public class LuceneToArffConverter {

    public String convert(Directory dir) throws Exception {
        int i,k;
        TreeMap<String, Integer> terminos = new TreeMap<String, Integer>();
    	IndexReader indice = IndexReader.open(dir);
        List<String> listaAtributos = new ArrayList<String>();
        TermEnum terms = indice.terms();
        Term t = null; //terms.term();
        int nterm = 0;
        do{
             t = terms.term();
             if(t!= null && terminos.get(t.text())==null){
                TermDocs td = indice.termDocs(t);
                int acum = 0;
                int ndoc = 0;
                int freq = 0;
                while(td.next()){
                   ndoc = td.doc();
                   Document doc = indice.document(ndoc);
                   if(doc.getFieldable("filename").stringValue().trim().equals(t.text().trim())){
                       acum = 0;
                       break;
                   }
                   freq = td.freq();
                   acum+=freq;
                }
                if(acum > 20){
                  terminos.put(t.text(), new Integer(nterm));
                  String texto = "@attribute "+"\""+t.text()+"\""+" numeric\n";
                  listaAtributos.add(texto);
                
                  nterm++;
                }
             }
            
             
        }while(terms.next());
        
        List<String> documentos = new ArrayList<String>();
        Set<String> setCategorias = new TreeSet<String>();
        for(i=0; i< indice.numDocs();i++){
            //System.out.print("{ ");
            TermFreqVector termFreqVector = indice.getTermFreqVector(i, "contents");
            String texto = " ";
            if (termFreqVector != null) {
                String[] term_usados = termFreqVector.getTerms();
                int[] freq = termFreqVector.getTermFrequencies();
                Iterator<String> nombresTerminos= terminos.keySet().iterator();
                while (nombresTerminos.hasNext()) {
                	String nombreTermino = nombresTerminos.next();
                    int frecuenciaTermino = 0;
                    for(k=0; (k< term_usados.length&& frecuenciaTermino == 0);k++){
                        String ter = term_usados[k];
                        if(nombreTermino.equals(ter)) {
                        	frecuenciaTermino = freq[k];
                        }
                    }
                    texto += " "+frecuenciaTermino+", ";
				}
            }
            Document docu = indice.document(i);
            texto += "'" + docu.getFieldable("categoria").stringValue() + "'\n";
            setCategorias.add(docu.getFieldable("categoria").stringValue());
            documentos.add(texto);
        }
        String archivo = "D:\\LuceneToArff.arff";
        FileOutputStream os = new FileOutputStream(archivo);
        String texto = "@relation LuceneToArff\n\n";
        byte[] contenido = texto.getBytes();
        os.write(contenido);
        System.out.println(texto);
        for (String atributo: listaAtributos){
        	texto = atributo;
            contenido = texto.getBytes();
            os.write(contenido);
            System.out.println(texto);
        }
        StringBuffer categorias = new StringBuffer("");
        for (String categoriaActual: setCategorias){
        	categorias.append("'" + categoriaActual + "'" + ",");
        }
        texto = "@attribute categoriaClasificacion " + "{" + categorias.toString() + "}";
        contenido = texto.getBytes();
        os.write(contenido);
        System.out.println(texto);
        texto = "\n\n@data\n\n";
        contenido = texto.getBytes();
        os.write(contenido);
        System.out.println(texto);
        for (String documento: documentos) {
        	texto = documento;
            contenido = texto.getBytes();
            os.write(contenido);
            System.out.println(texto);
        }
        os.close();
        indice.close();
        
        return archivo;
    }

}
